package raft

import (
	"bytes"
	"encoding/gob"
	"fmt"
	"log"
	"math/rand"
	"os"
	"raft-go/raft/storage"
	"sync"
	"time"
)

type NodeState int

const (
	Follower NodeState = iota
	Candidate
	Leader
	Dead
)

type CommitEntry struct {
	// 客户端提交的命令
	Command any
	// 客户端提交命令时的任期
	Term int
	// 客户端提交命令的索引
	Index int
}

const DebugCM = true

func (ns NodeState) String() string {
	switch ns {
	case Follower:
		return "Follower"
	case Candidate:
		return "Candidate"
	case Leader:
		return "Leader"
	case Dead:
		return "Dead"
	default:
		panic("Unreachable!")
	}
}

type LogEntry struct {
	Command any
	Term    int
}

type ConsensusModule struct {
	mu     sync.Mutex
	Id     int
	Peers  []int
	server *Server
	// 选取的轮次
	currentTerm int

	commitChan        chan<- CommitEntry
	commitReadyChan   chan struct{}
	commitReadyChanWg sync.WaitGroup
	triggerAEChan     chan struct{}
	commitIndex       int
	lastAppliedIndex  int
	// 记录各个Follower节点，下一条应该发送日志的下标
	nextIndex map[int]int
	// 记录是Leader节点复制日志到Follower节点时的日志下标
	matchIndex map[int]int

	log []LogEntry
	// 节点状态，包含Follower/Candidate/Leader/Dead四种状态
	state NodeState
	// 当前轮次投票给哪个节点
	votedFor           int
	electionResetEvent time.Time

	// 可持久化存储
	storage storage.Storage
}

func NewConsensusModule(id int, ids []int, s *Server, storage storage.Storage, ready <-chan any, commitChan chan<- CommitEntry) *ConsensusModule {
	cm := new(ConsensusModule)
	cm.Id = id
	cm.Peers = ids
	cm.server = s
	cm.state = Follower
	cm.currentTerm = 0
	cm.votedFor = -1
	cm.commitChan = commitChan
	cm.commitReadyChan = make(chan struct{}, 16)
	cm.commitIndex = -1
	cm.lastAppliedIndex = -1
	cm.nextIndex = make(map[int]int)
	cm.matchIndex = make(map[int]int)
	cm.triggerAEChan = make(chan struct{}, 1)
	cm.storage = storage

	if cm.storage.HasData() {
		cm.loadDataFromStorage(cm.storage)
	}

	go func() {
		<-ready
		cm.mu.Lock()
		cm.electionResetEvent = time.Now()
		cm.mu.Unlock()
		cm.runElectionTimer()
	}()

	cm.commitReadyChanWg.Add(1)
	go cm.CommitChansSender()

	return cm
}

func (cm *ConsensusModule) CommitChansSender() {
	defer cm.commitReadyChanWg.Done()

	for range cm.commitReadyChan {
		cm.mu.Lock()
		savedTerm := cm.currentTerm
		savedLastApplied := cm.lastAppliedIndex
		var entries []LogEntry
		if cm.commitIndex > cm.lastAppliedIndex {
			entries = cm.log[cm.lastAppliedIndex+1 : cm.commitIndex+1]
			cm.lastAppliedIndex = cm.commitIndex
		}
		cm.mu.Unlock()
		cm.Log("CommitChansSender entries=%v, savedLastApplied=%d", entries, savedLastApplied)

		for i, entry := range entries {
			cm.Log("send on commit chan i=%v, entry=%v", i, entry)
			cm.commitChan <- CommitEntry{
				Command: entry.Command,
				Term:    savedTerm,
				Index:   savedLastApplied + i + 1,
			}
		}
	}
	cm.Log("CommitChansSender done")
}

func (cm *ConsensusModule) Submit(command any) int {
	cm.mu.Lock()
	defer cm.mu.Unlock()
	cm.Log("submit received by %s, command=%v", cm.state.String(), command)
	if cm.state == Leader {
		submitIndex := len(cm.log)
		cm.log = append(cm.log, LogEntry{Command: command, Term: cm.currentTerm})
		cm.persistStorage()
		cm.Log("...log=%+v", cm.log)
		cm.triggerAEChan <- struct{}{}
		return submitIndex
	}
	return -1
}

func (cm *ConsensusModule) Log(format string, args ...any) {
	if DebugCM {
		format = fmt.Sprintf("[%d] ", cm.Id) + format
		log.Printf(format, args...)
	}
}

func (cm *ConsensusModule) Stop() {
	cm.mu.Lock()
	cm.state = Dead
	cm.Log("becomes Dead")
	cm.mu.Unlock()

	// 关闭本节点的消息通知管道，并等待所有消息提交协程
	close(cm.commitReadyChan)
	cm.commitReadyChanWg.Wait()
}

func (cm *ConsensusModule) Report() (id int, term int, isLeader bool) {
	cm.mu.Lock()
	defer cm.mu.Unlock()
	return cm.Id, cm.currentTerm, cm.state == Leader
}

func (cm *ConsensusModule) LastLogIndexAndTerm() (int, int) {
	if len(cm.log) > 0 {
		lastIndex := len(cm.log) - 1
		lastTerm := cm.log[lastIndex].Term
		return lastIndex, lastTerm
	}
	return -1, -1
}

// 随机生成一个选举超时时间，一般设定为150ms到300ms
func (cm *ConsensusModule) electionTimeout() time.Duration {
	if len(os.Getenv("RAFT_FORCE_MORE_REELECTION")) > 0 && rand.Intn(3) == 0 {
		return time.Duration(150) * time.Millisecond
	} else {
		return time.Duration(150+rand.Intn(150)) * time.Millisecond
	}
}

func (cm *ConsensusModule) runElectionTimer() {
	timeDuration := cm.electionTimeout()
	cm.mu.Lock()
	termStarted := cm.currentTerm
	cm.mu.Unlock()
	cm.Log("election timer started (%v), term=%d", timeDuration, termStarted)

	ticker := time.NewTicker(10 * time.Millisecond)
	defer ticker.Stop()
	for {
		<-ticker.C
		cm.mu.Lock()
		// 如果节点状态不是Candidate或Follower，那么不可进行选举（本身要么是Leader，要么节点下线）
		if cm.state != Candidate && cm.state != Follower {
			cm.Log("in election timer state = %s, bailing out", cm.state.String())
			cm.mu.Unlock()
			return
		}
		// 如果节点选举轮次未同步，即节点内容不一致
		if cm.currentTerm != termStarted {
			cm.Log("in election timer term changed, from %d to %d, bailing out", termStarted, cm.currentTerm)
			cm.mu.Unlock()
			return
		}

		// 如果在超时时间内，没有接收到Leader节点的心跳，或者没向别的节点投票，该节点就会发起一次选举
		if elapsed := time.Since(cm.electionResetEvent); elapsed >= timeDuration {
			cm.StartElection()
			cm.mu.Unlock()
			return
		}
		cm.mu.Unlock()
	}
}

func (cm *ConsensusModule) StartElection() {
	// 该节点发起选举，节点state首先变为Candidate，之后选举轮次+1
	// 并统计从其他节点的投票结果
	cm.state = Candidate
	cm.currentTerm += 1
	saveCurrentTerm := cm.currentTerm
	cm.electionResetEvent = time.Now()
	// 投票给自己[doge]
	cm.votedFor = cm.Id
	cm.Log("becomes Candidate(currentTerm=%d); log=%v", saveCurrentTerm, cm.log)

	votesReceived := 1
	for _, peer := range cm.Peers {
		go func(id int) {
			cm.mu.Lock()
			savedLastIndex, savedLastTerm := cm.LastLogIndexAndTerm()
			cm.mu.Unlock()
			args := RequestVoteArgs{
				Term:         saveCurrentTerm,
				CandidateId:  cm.Id,
				LastLogIndex: savedLastIndex,
				LastLogTerm:  savedLastTerm,
			}
			var reply RequestVoteReply
			cm.Log("send RequestVoteRequest to %d: %+v", id, args)
			err := cm.server.Call(id, "ConsensusModule.RequestVote", args, &reply)
			if err != nil {
				log.Printf("Peer Point is not Existed!\n")
				return
			}
			cm.mu.Lock()
			defer cm.mu.Unlock()
			cm.Log("received RequestVoteReply from %d, %+v", id, args)
			if cm.state != Candidate {
				cm.Log("state changed, state=[%s]", cm.state.String())
				return
			}

			// 返回的投票轮次大于当前投票轮次，说明已经选举出Leader了
			if reply.Term > saveCurrentTerm {
				cm.Log("RequestVote out of date!")
				cm.BecomFollower(reply.Term)
				return
			}
			// 获取票，对于2n+1节点系统，获取大于n个票，则成为Leader
			if reply.Term == saveCurrentTerm && reply.VoteGranted {
				votesReceived += 1
				if votesReceived*2 >= len(cm.Peers)+1 {
					cm.Log("wins election! Get %d votes", votesReceived)
					cm.BecomeLeader()
					return
				}
			}
		}(peer)
	}
	// 如果本次选举，没有选出Leader
	go cm.runElectionTimer()
}

func (cm *ConsensusModule) BecomFollower(term int) {
	cm.Log("becomes Follower with term=%d, log=%v", term, cm.log)
	cm.currentTerm = term
	cm.state = Follower
	cm.votedFor = -1
	cm.electionResetEvent = time.Now()
	go cm.runElectionTimer()
}

func (cm *ConsensusModule) BecomeLeader() {
	cm.state = Leader
	for _, peerId := range cm.Peers {
		cm.nextIndex[peerId] = len(cm.log)
		cm.matchIndex[peerId] = -1
	}
	cm.Log("become a Leader, term=%d, nextIndex=%v, matchIndex=%v, entries = %v", cm.currentTerm, cm.nextIndex, cm.matchIndex, cm.log)

	go func(heartBeatTimeout time.Duration) {
		cm.sendHeartBeat()
		timer := time.NewTimer(heartBeatTimeout)
		defer timer.Stop()
		for {
			doSend := false
			select {
			case <-timer.C:
				doSend = true
				timer.Stop()
				timer.Reset(heartBeatTimeout)
			case _, ok := <- cm.triggerAEChan:
				if ok {
					doSend = true
				} else {
					return
				}

				if !timer.Stop() {
					<-timer.C
				}
				timer.Reset(heartBeatTimeout)
			}

			if doSend {
				cm.mu.Lock()
				if cm.state != Leader {
					cm.mu.Unlock()
					return
				}
				cm.mu.Unlock()
				cm.sendHeartBeat()
			}
		}
	}(50 * time.Millisecond)
}

func (cm *ConsensusModule) sendHeartBeat() {
	cm.mu.Lock()
	// 如果该节点不是Leader节点，那么就无需发送心跳
	if cm.state != Leader {
		cm.mu.Unlock()
		return
	}
	savedCurrentTerm := cm.currentTerm
	cm.mu.Unlock()

	for _, peer := range cm.Peers {
		go func(id int) {
			cm.mu.Lock()
			next_index := cm.nextIndex[id]
			preLogIndex := next_index - 1
			preLogTerm := -1
			if preLogIndex >= 0 {
				preLogTerm = cm.log[preLogIndex].Term
			}
			entries := cm.log[next_index:]

			args := AppendEnriesArgs{
				Term:         savedCurrentTerm,
				LeaderId:     cm.Id,
				PrevLogIndex: preLogIndex,
				PrevLogTerm:  preLogTerm,
				Entries:      entries,
				LeaderCommit: cm.commitIndex,
			}
			cm.mu.Unlock()
			cm.Log("send AppendEntries to %d: next_index=%d, args=%+v", id, next_index, args)
			var reply AppendEntriesReply
			err := cm.server.Call(id, "ConsensusModule.AppendEntries", args, &reply)
			if err != nil {
				log.Printf("Peer Point is not Existed!\n")
				return
			}
			cm.mu.Lock()
			if reply.Term > cm.currentTerm {
				cm.Log("term out of date in heartbeat reply")
				cm.BecomFollower(reply.Term)
				cm.mu.Unlock()
				return
			}
			if cm.state == Leader && reply.Term == savedCurrentTerm {
				// 拒绝后，领导者对nextIndex进行减量，并重试AppendEntries RPC
				// 直到nextIndex达到领导者和跟随者日志匹配的点
				if !reply.Success {
					if reply.ConflictTerm >= 0 {
						lastIndexOfTerm := -1
						for i := len(cm.log) - 1; i >= 0; i-- {
							if cm.log[i].Term == reply.ConflictTerm {
								lastIndexOfTerm = i
								break
							}
						}
						if lastIndexOfTerm >= 0 {
							cm.nextIndex[id] = lastIndexOfTerm + 1
						} else {
							cm.nextIndex[id] = reply.ConflictIndex
						}
					} else {
						cm.nextIndex[id] = reply.ConflictIndex
					}
					cm.Log("AppendEntries reply from %d !success: nextIndex := %d", id, next_index-1)
					cm.mu.Unlock()
					return
				}
				// if success
				cm.nextIndex[id] = next_index + len(entries)
				cm.matchIndex[id] = cm.nextIndex[id] - 1
				savedCommitIndex := cm.commitIndex
				for i := cm.commitIndex + 1; i < len(cm.log); i++ {
					if cm.log[i].Term == cm.currentTerm {
						matchedCount := 1
						for _, p := range cm.Peers {
							if cm.matchIndex[p] >= i {
								matchedCount += 1
							}
						}
						// 大部分节点已经同步，更新提交的日志索引
						if matchedCount*2 > len(cm.Peers)+1 {
							cm.commitIndex = i
						}
					}
				}
				cm.Log("AppendEntries received from %d at term=%d [matchedIndex=%v, nextIndex=%v]", id, savedCurrentTerm, cm.matchIndex, cm.nextIndex)
				// 如果可以提交（即获得了大部分节点的认可），则向通道发送信息，提交日志条目
				if savedCommitIndex != cm.commitIndex {
					cm.Log("leader set commit index = %d", cm.commitIndex)
					cm.mu.Unlock()
					cm.commitReadyChan <- struct{}{}
					cm.triggerAEChan <- struct{}{}
				} else {
					cm.mu.Unlock()
				}
			} else {
				cm.mu.Unlock()
			}
		}(peer)
	}
}

// RPC implement
func (cm *ConsensusModule) RequestVote(args RequestVoteArgs, reply *RequestVoteReply) error {
	cm.mu.Lock()
	defer cm.mu.Unlock()
	if cm.state == Dead {
		return nil
	}

	lastLogIndex, lastTerm := cm.LastLogIndexAndTerm()
	cm.Log("RequestVote: %+v [currentTerm=%d, voteFor=%d, log index/temr=(%d, %d)]", args, cm.currentTerm, cm.votedFor, lastLogIndex, lastTerm)

	if args.Term > cm.currentTerm {
		cm.Log("...term out of date in RequestVote")
		cm.BecomFollower(args.Term)
	}

	/*
		如何判断B能够给A投票的（A->B发送请求）
		1. 首先A和B之间的任期一致，如果任期不一致，选票不合理，不应被采纳
		2. 其次，B节点并未投票（即可以投票），或者B节点在当前任期中已经投票给A了
		3. 最后还需判断节点之间的日志状态是否一致，如果A节点的日志落后于（index，term）B，
		   那么B就不会给A投票
	*/
	if args.Term == cm.currentTerm &&
		(cm.votedFor == -1 || cm.votedFor == args.CandidateId) &&
		(args.LastLogTerm > lastTerm || (args.LastLogTerm == lastTerm && args.LastLogIndex >= lastLogIndex)) {
		cm.votedFor = args.CandidateId
		cm.electionResetEvent = time.Now()
		reply.VoteGranted = true
	} else {
		reply.VoteGranted = false
	}
	reply.Term = cm.currentTerm
	cm.persistStorage()
	cm.Log("...RequestVote reply: %+v", reply)
	return nil
}

func (cm *ConsensusModule) AppendEntries(args AppendEnriesArgs, reply *AppendEntriesReply) error {
	cm.mu.Lock()
	defer cm.mu.Unlock()
	if cm.state == Dead {
		return nil
	}

	cm.Log("Append Log entry = %+v", args)

	// 判断任期是否不一致
	if args.Term > cm.currentTerm {
		cm.Log("...term is out of date")
		cm.BecomFollower(args.Term)
	}

	reply.Success = false
	if args.Term == cm.currentTerm {
		if cm.state != Follower {
			cm.Log("becomes Follower...")
			cm.BecomFollower(args.Term)
		}
		cm.electionResetEvent = time.Now()
		// 如果是一次未提交的请求，那么本节点也同步状态
		if args.PrevLogIndex == -1 ||
			(args.PrevLogTerm == cm.log[args.PrevLogIndex].Term && args.PrevLogIndex < len(cm.log)) {
			reply.Success = true
			logInsertIndex := args.PrevLogIndex + 1
			entriesIndex := 0
			// 对齐日志以及节点之间的任期时钟
			for {
				if logInsertIndex >= len(cm.log) || entriesIndex >= len(args.Entries) {
					break
				}
				if cm.log[logInsertIndex].Term != args.Entries[entriesIndex].Term {
					break
				}
				logInsertIndex++
				entriesIndex++
			}
			// 添加entries到节点日志序列中
			if entriesIndex < len(args.Entries) {
				cm.Log("...append log entries %v from index=%d", args.Entries[entriesIndex:], logInsertIndex)
				cm.log = append(cm.log[:logInsertIndex], args.Entries[entriesIndex:]...)
				cm.Log("...log now is %v", cm.log)
			}
			// 如果Leader节点提交了条目则Follower节点也提交对应日志条目
			if args.LeaderCommit > cm.commitIndex {
				cm.commitIndex = min(args.LeaderCommit, len(cm.log)-1)
				cm.Log("... setting commitIndex=%d", cm.commitIndex)
				cm.commitReadyChan <- struct{}{}
			}
		} else {
			if args.PrevLogIndex >= len(cm.log) {
				reply.ConflictIndex = len(cm.log)
				reply.ConflictTerm = -1
			} else {
				reply.ConflictTerm = cm.log[args.PrevLogIndex].Term

				var i int
				for i = args.PrevLogIndex - 1; i >= 0; i-- {
					if cm.log[i].Term != reply.ConflictTerm {
						break
					}
				}
				reply.ConflictIndex = i + 1
			}
		}
	}
	reply.Term = cm.currentTerm
	cm.persistStorage()
	cm.Log("AppendEntries reply is %+v", *reply)

	return nil
}

func (cm *ConsensusModule) loadDataFromStorage(s storage.Storage) {
	if data, found := s.Get("currentTerm"); !found {
		log.Fatalf("current term is not found")
	} else {
		err := gob.NewDecoder(bytes.NewBuffer(data)).Decode(&cm.currentTerm)
		if err != nil {
			log.Fatalf("decode current term failed, err=%s\n", err.Error())
		}
	}

	if data, found := s.Get("votedFor"); !found {
		log.Fatalf("votedFor is not found")
	} else {
		err := gob.NewDecoder(bytes.NewBuffer(data)).Decode(&cm.votedFor)
		if err != nil {
			log.Fatalf("decode votedFor failed, err=%s\n", err.Error())
		}
	}

	if data, found := s.Get("log"); !found {
		log.Fatalf("log is not found")
	} else {
		err := gob.NewDecoder(bytes.NewBuffer(data)).Decode(&cm.log)
		if err != nil {
			log.Fatalf("decode log failed, err=%s\n", err.Error())
		}
	}
}

// 可持久化存储节点的currentTerm, votedFor以及log
// cm is Locked
func (cm *ConsensusModule) persistStorage() {
	var termEncoder bytes.Buffer
	if err := gob.NewEncoder(&termEncoder).Encode(cm.currentTerm); err != nil {
		log.Fatalf("encode current term failed, err=%s", err.Error())
	}
	cm.storage.Set("currentTerm", termEncoder.Bytes())

	var votedForEncoder bytes.Buffer
	if err := gob.NewEncoder(&votedForEncoder).Encode(cm.votedFor); err != nil {
		log.Fatalf("encode voted for failed, err=%s", err.Error())
	}
	cm.storage.Set("votedFor", votedForEncoder.Bytes())

	var logEncoder bytes.Buffer
	if err := gob.NewEncoder(&logEncoder).Encode(cm.log); err != nil {
		log.Fatalf("encode log failed, err=%s", err.Error())
	}
	cm.storage.Set("log", logEncoder.Bytes())
}
